"""
MongoDB存储实现
用于存储金融新闻数据、实体数据等非结构化数据
"""
import pymongo
from typing import Dict, List, Any
from pymongo.errors import PyMongoError
from fin_senti_entity_platform.data_collection.storage.base_storage import BaseStorage, StorageFactory
from fin_senti_entity_platform.utils.logger import Logger


class MongoStorage(BaseStorage):
    """MongoDB存储实现类"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化MongoDB存储
        
        Args:
            config: MongoDB配置参数
        """
        super().__init__(config)
        self.logger = Logger().get_logger(__name__)
        self.client = None
        self.db = None
        self._connect()
        
    def _connect(self):
        """\连接MongoDB数据库"""
        try:
            # 创建MongoDB连接字符串
            host = self.config.get('host', 'localhost')
            port = self.config.get('port', 27017)
            username = self.config.get('username', None)
            password = self.config.get('password', None)
            db_name = self.config.get('db_name', 'financial_data')
            
            # 构建连接URI
            uri_parts = [f'mongodb://']
            if username and password:
                uri_parts.append(f'{username}:{password}@')
            uri_parts.append(f'{host}:{port}/')
            if username:
                uri_parts.append(f'?authSource=admin')
                
            uri = ''.join(uri_parts)
            
            # 连接MongoDB
            self.client = pymongo.MongoClient(
                uri,
                serverSelectionTimeoutMS=self.config.get('timeout', 5000),
                maxPoolSize=self.config.get('max_pool_size', 100)
            )
            
            # 选择数据库
            self.db = self.client[db_name]
            
            # 验证连接
            self.client.server_info()
            self.logger.info(f"成功连接到MongoDB数据库: {db_name}")
            
        except PyMongoError as e:
            self.logger.error(f"MongoDB连接失败: {str(e)}")
            raise
    
    def save(self, data: Dict[str, Any], collection_name: str, **kwargs) -> bool:
        """
        保存单条数据到MongoDB集合
        
        Args:
            data: 要保存的数据（字典格式）
            collection_name: 集合名称
            **kwargs: 其他参数
                - upsert: 是否更新插入
                - filter: 更新条件（用于upsert）
        
        Returns:
            bool: 保存是否成功
        """
        try:
            collection = self.db[collection_name]
            
            if kwargs.get('upsert', False) and 'filter' in kwargs:
                # 使用upsert模式
                result = collection.update_one(
                    kwargs['filter'],
                    {'$set': data},
                    upsert=True
                )
                return result.acknowledged
            else:
                # 直接插入
                result = collection.insert_one(data)
                return result.acknowledged
                
        except PyMongoError as e:
            self.logger.error(f"保存数据到集合 {collection_name} 失败: {str(e)}")
            return False
    
    def batch_save(self, data_list: List[Dict[str, Any]], collection_name: str, **kwargs) -> bool:
        """
        批量保存数据到MongoDB集合
        
        Args:
            data_list: 数据列表
            collection_name: 集合名称
            **kwargs: 其他参数
        
        Returns:
            bool: 批量保存是否成功
        """
        try:
            if not data_list:
                return True
                
            collection = self.db[collection_name]
            result = collection.insert_many(data_list)
            
            if len(result.inserted_ids) == len(data_list):
                self.logger.info(f"成功批量保存 {len(data_list)} 条数据到集合 {collection_name}")
                return True
            else:
                self.logger.warning(f"批量保存数据不完整，预期 {len(data_list)} 条，实际保存 {len(result.inserted_ids)} 条")
                return False
                
        except PyMongoError as e:
            self.logger.error(f"批量保存数据到集合 {collection_name} 失败: {str(e)}")
            return False
    
    def query(self, collection_name: str, query_params: Dict[str, Any], **kwargs) -> List[Dict[str, Any]]:
        """
        查询MongoDB集合中的数据
        
        Args:
            collection_name: 集合名称
            query_params: 查询条件
            **kwargs: 其他参数
                - projection: 投影字段
                - sort: 排序条件
                - limit: 返回记录数限制
                - skip: 跳过记录数
        
        Returns:
            List[Dict[str, Any]]: 查询结果列表
        """
        try:
            collection = self.db[collection_name]
            
            # 构建查询
            query = collection.find(query_params)
            
            # 应用投影
            if 'projection' in kwargs:
                query = query.projection(kwargs['projection'])
                
            # 应用排序
            if 'sort' in kwargs:
                query = query.sort(kwargs['sort'])
                
            # 应用分页
            if 'skip' in kwargs:
                query = query.skip(kwargs['skip'])
                
            if 'limit' in kwargs:
                query = query.limit(kwargs['limit'])
                
            # 执行查询并返回结果
            results = list(query)
            self.logger.info(f"查询集合 {collection_name}，返回 {len(results)} 条记录")
            return results
            
        except PyMongoError as e:
            self.logger.error(f"查询集合 {collection_name} 失败: {str(e)}")
            return []
    
    def update(self, collection_name: str, query_params: Dict[str, Any], update_data: Dict[str, Any], **kwargs) -> bool:
        """
        更新MongoDB集合中的数据
        
        Args:
            collection_name: 集合名称
            query_params: 查询条件
            update_data: 更新数据
            **kwargs: 其他参数
                - multi: 是否更新多条记录
        
        Returns:
            bool: 更新是否成功
        """
        try:
            collection = self.db[collection_name]
            
            if kwargs.get('multi', False):
                result = collection.update_many(query_params, {'$set': update_data})
                self.logger.info(f"更新集合 {collection_name} 中的 {result.modified_count} 条记录")
            else:
                result = collection.update_one(query_params, {'$set': update_data})
                self.logger.info(f"更新集合 {collection_name} 中的 {result.modified_count} 条记录")
                
            return result.acknowledged
            
        except PyMongoError as e:
            self.logger.error(f"更新集合 {collection_name} 失败: {str(e)}")
            return False
    
    def delete(self, collection_name: str, query_params: Dict[str, Any], **kwargs) -> bool:
        """
        删除MongoDB集合中的数据
        
        Args:
            collection_name: 集合名称
            query_params: 删除条件
            **kwargs: 其他参数
                - multi: 是否删除多条记录
        
        Returns:
            bool: 删除是否成功
        """
        try:
            collection = self.db[collection_name]
            
            if kwargs.get('multi', False):
                result = collection.delete_many(query_params)
                self.logger.info(f"删除集合 {collection_name} 中的 {result.deleted_count} 条记录")
            else:
                result = collection.delete_one(query_params)
                self.logger.info(f"删除集合 {collection_name} 中的 {result.deleted_count} 条记录")
                
            return result.acknowledged
            
        except PyMongoError as e:
            self.logger.error(f"删除集合 {collection_name} 失败: {str(e)}")
            return False
    
    def close(self):
        """关闭MongoDB连接"""
        if self.client:
            self.client.close()
            self.logger.info("MongoDB连接已关闭")


# 注册MongoDB存储实现
StorageFactory.register_storage('mongodb', MongoStorage)